# 双偏振雷达基数据，标准格式，CINRAD/SAD，读取类库

import numpy as np
import pandas as pd
import operator

from CinradSADProtocol import dtype_SAD  # 数据结构
# 引入工具函数
from util import _unpack_from_buf


class CinradSADData(object):
    """
    读取CINRAD-SAD基数据文件，张深寿，2021-10-13。此类 函数 util 代码来自 郑玉的pycwr
    """

    # 构造类
    def __init__(self, filename, station_lon=None, station_lat=None, station_alt=None):
        super(CinradSADData, self).__init__()
        self.filename = filename  # 雷达基数据文件名

        # 数据结构头
        self.GenericHeader = []  # 通用头块
        self.SiteConfig = []  # 站点配置块
        self.TaskConfig = []  # 扫描任务配置块

        self.CutConfigs = []  # 各层PPI CUT的配置块

        self.radials = []  # 所有径向数据

        with open(filename, 'rb') as f:  # 打开文件
            struct_buf = f.read(32)  # 读取通用头块
            self.GenericHeader, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.GenericHeader())

            struct_buf = f.read(128)  # 读站点配置块
            self.SiteConfig, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.SiteConfig())

            struct_buf = f.read(256)  # 读扫描任务配置块
            self.TaskConfig, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.TaskConfig())

            for i in range(self.TaskConfig["CutNumber"]):  # 读各层PPI CUT的配置块
                struct_buf = f.read(256)
                CutConfig, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.CutConfig())
                self.CutConfigs.append(CutConfig)

                # 读取所有径向数据
            while True:
                # 读取径向
                struct_buf = f.read(64)  # 读径向头块
                if len(struct_buf) < 64:  # 从文件读取到的小于64字节，说明读取不到新的数据了，文件读取完毕，跳出
                    break

                Radial = []
                # 读取径向数据结构信息
                RadialHeader, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.RadialHeader())
                Radial.append(RadialHeader)

                # 读取每个数据类型的数据，如 dBT、dBZ、V、W、ZDR、CC、KDP、PhidP等
                for i in range(RadialHeader["MomentNumber"]):  # RadialHeader["MomentNumber"]  指示该径向有几种数据
                    MomentData = {}  # 存储某一类型的数据

                    struct_buf = f.read(32)  # 读某一类型数据的数据头
                    RadialDataHeader, size_tmp = _unpack_from_buf(struct_buf, 0, dtype_SAD.RadialDataHeader())
                    MomentData.update(RadialDataHeader)

                    # 读取数据，RadialDataHeader["Length"] 指示了数据长度
                    dataBuf = f.read(RadialDataHeader["Length"])

                    if RadialDataHeader["BinLength"] == 2:
                        # 转换并追加到一根径向的所有类型数据集中
                        Value = np.frombuffer(dataBuf, dtype="u2")
                        MomentData['Data'] = Value
                    else:
                        # 转换并追加到一根径向的所有类型数据集中
                        Value = np.frombuffer(dataBuf, dtype="u1")
                        MomentData['Data'] = Value

                    Radial.append(MomentData)

                # 径向数据追加到径向数据集中
                self.radials.append(Radial)

        # 所有径向排序，按仰角号排序，再按方位排序（从0°开始），排序后显示图像会少一根径向，不知道为啥？：(
        # self.radials = sorted(self.radials, key = lambda radial : (radial[0]["ElevationNumber"], radial[0]["Azimuth"]) )

    def GetPPIData(self, types, Ele=1, MaxR=230, minValue=-100.0, MaxValue=100.0):
        # 得到某一层次、某一类型的PPI数据
        # 参数：types 数据类型数组，可以多种数据，详细见 CinradSADProtocol.py 里的 表2-6 定义描述
        #          Ele  仰角层次
        #          MaxR 需要获取数据的最大距离
        #          minValue 需要过滤数据的最小值
        #          MaxValue 需要过滤数据的最大值

        PPI_Datas = [radial for radial in self.radials if round(radial[0]["ElevationNumber"], 1) == Ele]  # 得到一个层次的PPI数据

        # 径向方位
        AzAll = np.array([iradial[0]['Azimuth'] for iradial in PPI_Datas])  # 一个PPI所有径向的方位集合
        AzCount = len(AzAll)  # 一个PPI总径向数

        # 分辨率
        if self.CutConfigs[Ele]["LogResolution"] > 0:  # 如果有反射率因子，用反射率因子的分辨率
            Reso = self.CutConfigs[Ele]["LogResolution"] / 1000  # 距离分辨率，单位：km
        else:  # 否则用多普勒速度的分辨率
            Reso = self.CutConfigs[Ele]["DopplerResolution"] / 1000  # 距离分辨率，单位：km

        # 准备数据
        Ref_RCount = int(MaxR / Reso)  # 距离径向数据个数

        Ref_GateSize = Reso  # 库长，单位：km
        RangeToFirstGateOfRef = 0  # 第一个库的距离，单位：km
        Ref_R = np.linspace(RangeToFirstGateOfRef, (Ref_RCount - 1) * Ref_GateSize, Ref_RCount)  # 距离

        # 构建极坐标的方位距离网格
        Ref_Ri, Ref_Azi = np.meshgrid(Ref_R, AzAll)

        # 极坐标转换后的散点 X、Y、H坐标，初始化
        # todo: AzCount径向数 Ref_RCount 距离径向数据个数
        Ref_X = np.full((AzCount, Ref_RCount), np.nan)
        Ref_Y = np.full((AzCount, Ref_RCount), np.nan)
        Ref_H = np.full((AzCount, Ref_RCount), np.nan)

        # 用三角函数转换极坐标为笛卡尔坐标
        Ref_Y = Ref_Ri * np.sin(np.radians(Ref_Azi))
        Ref_X = Ref_Ri * np.cos(np.radians(Ref_Azi))

        # 得到该层对应的所有仰角
        Eles = np.full((AzCount, Ref_RCount), np.nan)
        for i in range(AzCount):
            EleValue = PPI_Datas[i][0]["Elevation"]
            eleRadial = np.full(Ref_RCount, EleValue)  # 一个径向所有的相同仰角
            eleRadial = np.sin(np.radians(eleRadial))  # 转换成sin值，供测高公式后续使用
            Eles[i] = eleRadial

        H0 = self.SiteConfig["AntennaHeight"] / 1000.  # 1.507 #雷达站海拔，单位：km
        Ref_H = H0 + Ref_Ri * Eles + Ref_Ri * Ref_Ri / 17000  # 雷达测高公式，计算 高度

        # 得到径向数据的所含有的数据类型代码
        number = PPI_Datas[0][0]['MomentNumber']
        Moments = []
        for k in range(1, number + 1):
            Moments.append(PPI_Datas[0][k]['DataType'])

        # 需要返回的数据集
        MomentDatas = []
        for typesIndex in range(len(types)):  # 每个需要的数据类型
            # 找到所需数据类型的在径向的数据数组索引，如果没有数据为 -1
            indexNumber = Moments.index(types[typesIndex]) + 1 if types[typesIndex] in Moments else -1

            # 初始化数据为空
            DataValues = np.full((AzCount, Ref_RCount), np.nan)

            if indexNumber > -1:  # 如果有数据
                # 检查数据是否有超过数据长度，超过则用数据里的长度
                RCount = Ref_RCount
                if RCount > PPI_Datas[0][indexNumber]["Length"]:  # 每个径向数据最大个数
                    RCount = PPI_Datas[0][indexNumber]["Length"]

                # 计算数据的Offset和Scale
                offset = PPI_Datas[0][indexNumber]['Offset']
                Scale = PPI_Datas[0][indexNumber]['Scale']
                # 每个方位，获取相应数据
                for i in range(AzCount):
                    DataValuesTemp = np.where(PPI_Datas[i][indexNumber]['Data'] > 4,
                                              (PPI_Datas[i][indexNumber]['Data'].astype(int) - offset) / Scale,
                                              np.nan)  # 计算值
                    DataValuesTemp = DataValuesTemp[0: RCount]

                    DataValuesTemp = np.where(DataValuesTemp < minValue, np.nan, DataValuesTemp)  # 根据要求，把大于最大小于最小值删除
                    DataValuesTemp = np.where(DataValuesTemp > MaxValue, np.nan, DataValuesTemp)

                    DataValues[i] = DataValuesTemp  # 取数据区间

            # 加入数据集，如果没有数据加入的是初始化空值
            MomentDatas.append(DataValues)

        return MomentDatas, Ref_X, Ref_Y, Ref_H